import torch.nn as nn
import numpy as np
import torch


class Aoct_Model(nn.Module):
    def __init__(self, num_classes=1):
        super().__init__()
        self.layers = self.get_model([1, 16, 32, 48, 64, 64])
        self.avgpool = nn.AdaptiveAvgPool2d((1,1))
        self.maxpool = nn.AdaptiveMaxPool2d((1,1))
        self.leaner0 = nn.Linear(in_features=4224, out_features=100)
        self.leaner1 = nn.Linear(in_features=100, out_features=num_classes)

        self.dropout1 = nn.Dropout(p=0.33)  # 添加Dropout层
        self.dropout2 = nn.Dropout(p=0.33)  # 添加Dropout层

    def get_model(self, channels):
        layers = nn.ModuleList()
        for i in range((len(channels) - 1)):
            stride = 2
            layers.append(nn.Conv2d(in_channels=channels[i], out_channels=channels[i + 1], kernel_size=3, padding=1, stride=stride))
            layers.append(nn.BatchNorm2d(channels[i + 1]))
            layers.append(nn.PReLU())
        return layers

    def forward(self, x):
        for layer in self.layers:
            x = layer(x)
        pool_avg = self.avgpool(x)[:,:,0,0]
        pool_max = self.maxpool(x)[:,:,0,0]
        x = torch.flatten(x, 1)
        output = torch.cat((x, pool_avg, pool_max), dim=1)

        x = self.dropout1(output)  # 应用第二个Dropout层
        x = self.leaner0(x)
        x = self.dropout2(x)  # 应用第二个Dropout层
        x = self.leaner1(x)
        return x

if __name__ == '__main__':
    model = Aoct_Model()
    device = torch.device('cuda')
    model = model.to(device=device)
    x = torch.rand((1, 1, 256, 256))
    x = x.to(device=device)
    result = model.forward(x)
